# ocr_processor.py
import os
from posixpath import dirname
from paddleocr import PaddleOCR
import re
import cv2
import numpy as np
import json


# 12：7
def standardize_image(image_path, target_size=(1206, 805)):
    """
    将发票图像标准化为统一尺寸（建议：A4 300dpi ≈ 2480x3508，可等比缩放）
    """
    img = cv2.imread(image_path)
    img = cv2.resize(img, target_size)  # 如 (1240, 1754) 为 1/2 A4 尺寸
    return img


class InvoiceOCR:
    def __init__(self):
        self.ocr = PaddleOCR(
            use_angle_cls=True,
            lang="ch",
            det_limit_side_len=960,
            rec_batch_num=16,
            # use_mp=False  # 避免多进程问题
        )

    def extract_text(self, image_path):
        result = self.ocr.ocr(image_path, cls=True)
        # 提取所有文本行
        texts = []
        for line in result:
            for word in line:
                text = word[1][0]  # (bbox, (text, score))
                texts.append(text.strip())
        return texts


# ocr_processor.py
class InvoiceExtractor:
    def __init__(self, texts):
        self.texts = texts
        self.result = self.init_structure()

    def init_structure(self):
        return {
            "invoice_type": "",
            "invoice_number": "",
            "issue_date": "",
            "buyer": {"name": "", "code": ""},
            "seller": {"name": "", "code": ""},
            "items": [
                {
                    "name": "",
                    "spec": "",
                    "unit": "",
                    "num": "",
                    "price": "",
                    "amount": "",
                    "rate": "",
                    "rateAmount": "",
                }
            ],
            "total": {
                "amount_total": "",
                "tax_total": "",
                "total_in_words": "",
                "total_in_digits": "",
            },
            "remark": "",
            "open": "",
        }

    def extract(self):
        self.extract_invoice_header()
        self.extract_invoice_info()
        self.extract_parties()
        self.extract_items_and_total()
        return self.result

    def extract_invoice_header(self):
        # 发票名称
        for text in self.texts:
            if "增值税专用发票" in text:
                self.result["invoice_name"] = text
                self.result["invoice_type"] = "专用发票"
                break

    def extract_invoice_info(self):
        # 发票代码、号码、日期
        for i, text in enumerate(self.texts):
            if "发票代码" in text and i + 1 < len(self.texts):
                self.result["invoice_code"] = self.texts[i + 1]
            elif "发票号码" in text and i + 1 < len(self.texts):
                self.result["invoice_number"] = self.texts[i + 1]
            elif "开票日期" in text and i + 1 < len(self.texts):
                self.result["issue_date"] = self.texts[i + 1]

    def extract_cipher_text(self):
        # 密码区通常是一长串特殊字符
        for text in self.texts:
            if re.match(r"^[0-9A-Za-z/<>*+\-]+$", text) and len(text) > 50:
                self.result["cipher_text"] = text
                break

    def extract_parties(self):
        # 购买方、销售方信息
        sections = {"buyer": ["购买方", "购货单位"], "seller": ["销售方", "销货单位"]}
        for key, keywords in sections.items():
            for keyword in keywords:
                for i, text in enumerate(self.texts):
                    if keyword in text:
                        # 名称
                        if i + 1 < len(self.texts):
                            self.result[key]["name"] = self.texts[i + 1]
                        # 纳税人识别号
                        if i + 2 < len(self.texts):
                            self.result[key]["tax_id"] = self.texts[i + 2]
                        # 地址电话
                        if i + 3 < len(self.texts):
                            self.result[key]["address_phone"] = self.texts[i + 3]
                        # 开户行及账号
                        if i + 4 < len(self.texts):
                            self.result[key]["bank_account"] = self.texts[i + 4]
                        break

    def extract_items_and_total(self):
        # 货物或应税劳务、服务
        in_items = False
        for text in self.texts:
            if "货物或应税劳务、服务" in text:
                in_items = True
                continue
            if in_items and "合计" in text:
                break
            if in_items and len(text) > 2:
                # 尝试提取名称、金额、税率、税额
                item = {"name": text, "amount": "", "tax_rate": "", "tax_amount": ""}
                # 后续可用正则匹配金额和税率
                self.result["items"].append(item)

        # 合计金额、税额、价税合计
        for i, text in enumerate(self.texts):
            if "价税合计" in text:
                if i + 1 < len(self.texts):
                    self.result["total"]["total_in_words"] = self.texts[i + 1]
                if i + 2 < len(self.texts):
                    self.result["total"]["total_in_digits"] = self.texts[i + 2]
            elif "金额合计" in text and i + 1 < len(self.texts):
                self.result["total"]["amount_total"] = self.texts[i + 1]
            elif "税额合计" in text and i + 1 < len(self.texts):
                self.result["total"]["tax_total"] = self.texts[i + 1]

        # 提取开票人、复核、收款人（通常在右下角）
        for i, text in enumerate(self.texts):
            if "收款人" in text and i + 1 < len(self.texts):
                self.result["payee"] = self.texts[i + 1]
            elif "复核" in text and i + 1 < len(self.texts):
                self.result["reviewer"] = self.texts[i + 1]
            elif "开票人" in text and i + 1 < len(self.texts):
                self.result["issuer"] = self.texts[i + 1]


# ocr_processor.py
class TemplateOCR:
    def __init__(self, template_dir="templates"):
        # 如果未传路径，则使用与当前文件同级的 templates/vat_invoice.json
        # if template_path is None:
        #     current_dir = os.path.dirname(__file__)  # 获取当前文件所在目录
        #     # 普通发票
        #     template_path = os.path.join(current_dir, "templates", "vat_invoice.json")

        # # 检查文件是否存在
        # if not os.path.exists(template_path):
        #     raise FileNotFoundError(f"模板文件未找到: {template_path}")

        # with open(template_path, "r", encoding="utf-8") as f:
        #     self.template = json.load(f)
        # self.ocr = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=False)
        self.template_dir = template_dir
        self.ocr = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=False)

    def recognize_region(self, img, x, y, w, h):
        """
        裁剪区域并 OCR 识别，支持多行文本合并
        """
        # 确保坐标和尺寸为整数
        x, y, w, h = int(x), int(y), int(w), int(h)
        # 边界检查
        # h_img, w_img = img.shape[:2]
        # x = max(0, x)
        # y = max(0, y)
        # w = min(w, w_img - x)
        # h = min(h, h_img - y)
        # if w <= 0 or h <= 0:
        #     return ""

        crop = img[y:h, x:w]
        # print(f"Crop:{crop}...x:{x},w:{w},y:{y},h:{h}")
        result = self.ocr.ocr(crop, cls=True)
        texts = []
        for line in result:
            for word_info in line:
                word = word_info[1][0]  # 提取文本
                texts.append(word.strip())
        return " ".join(texts)

    def identify_invoice_type(self, img):
        # 假设抬头位于固定的区域，这里只是一个例子，需要根据实际情况调整
        head_area = {"x": 220, "y": 54, "w": 879, "h": 117}
        text = self.recognize_region(
            img, head_area["x"], head_area["y"], head_area["w"], head_area["h"]
        )
        print(f'抬头识别结果为：{text}')
        # 根据文本内容判断发票类型
        if "电子发票" in text:
            return "vat_invoice.json"
        # elif "增值税专用发票" in text:
        # return "electronic_invoice.json"
        else:
            return "default_invoice.json"

    def load_template(self, template_name):
        current_dir = os.path.dirname(__file__)  # 获取当前文件所在目录

        with open(
            os.path.join(current_dir, self.template_dir, template_name),
            "r",
            encoding="utf-8",
        ) as f:
            return json.load(f)

    def extract_from_image(self, image_path):
        # 1. 标准化图像
        img = standardize_image(image_path, target_size=(1206, 805))

        # 1.5 获取发票抬头，初始化截取模板
        invoice_type = self.identify_invoice_type(img)
        template = self.load_template(invoice_type)

        # 2. 递归提取所有字段
        result = self._extract_recursive(template, img)

        # 3. 结构化输出（可选：映射到标准结构）
        structured = self.map_to_structure(result)
        return structured

    def _extract_recursive(self, template, img):
        """
        递归提取模板中的所有文本
        """
        result = {}
        for key, value in template.items():
            if isinstance(value, dict):
                if "x" in value and "y" in value and "w" in value and "h" in value:
                    # 是叶子节点（有坐标)
                    x, y, w, h = value["x"], value["y"], value["w"], value["h"]
                    text = self.recognize_region(img, x, y, w, h)
                    print(f"text:{text}")

                    if "名称：" in text:
                        text = text.replace("名称：", "")
                    if "统一社会信用代码/纳税人识别号：" in text:
                        text = text.replace("统一社会信用代码/纳税人识别号：", "")

                    result[key] = text
                else:
                    # 是中间节点（嵌套结构）
                    zj = self._extract_recursive(value, img)
                    print(f"Key:{key},中间节点:{zj}")
                    result[key] = zj
            else:
                # 兜底：非 dict 直接跳过（理论上不会出现）
                result[key] = str(value)
        return result

    def map_to_structure(self, raw):
        """
        将 raw 字段映射为最终结构
        """
        return raw
        # return {
        #     "invoice_type": raw.get("invoice_type", ""),
        #     "invoice_number": raw.get("invoice_number", ""),
        #     "issue_date": raw.get("issue_date", ""),
        #     "buyer": {
        #         "name": raw.get("name", ""),
        #         "code": raw.get("code", ""),
        #     },
        #     "seller": {
        #         "name": raw.get("name", ""),
        #         "code": raw.get("code", ""),
        #     },
        #     "items": [
        #         {
        #             "name": raw.get("name", ""),
        #             "spec": raw.get("spec", ""),
        #             "unit": raw.get("unit", ""),
        #             "num": raw.get("num", ""),
        #             "price": raw.get("price", ""),
        #             "amount": raw.get("amount", ""),
        #             "rate": raw.get("rate", ""),
        #             "rateAmount": raw.get("rateAmount", ""),
        #         }
        #     ],
        #     "total": {
        #         "amount_total": raw.get("total_amount", ""),
        #         "tax_total": raw.get("total_tax", ""),
        #         "total_in_words": raw.get("total_in_words", ""),
        #         "total_in_digits": raw.get("total_in_digits", ""),
        #     },
        #     "remark": raw.get("remark", ""),
        #     "open": raw.get("open", ""),
        # }
